import backtrader as bt
import datetime
from data_feed import etf_daily, MyStockData
class MovingAverageTurnStrategy(bt.Strategy):
    params = (
        ('ma_period', 5),
    )

    def __init__(self):
        self.ma = bt.indicators.MovingAverageSimple(self.data.close, period=self.params.ma_period)
        self.order = None

    def log(self, txt, dt=None):
        dt = dt or self.datas[0].datetime.date(0)
        print(f'{dt.isoformat()}, {txt}')

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f'BUY EXECUTED, Price: {order.executed.price}, Cost: {order.executed.value}, Comm: {order.executed.comm}')
            elif order.issell():
                self.log(f'SELL EXECUTED, Price: {order.executed.price}, Cost: {order.executed.value}, Comm: {order.executed.comm}')

        self.order = None

    def next(self):
        if self.order:
            return

        if not self.position:
            if self.ma[0] > self.ma[-1] and self.ma[-1] < self.ma[-2]:
                self.log(f'BUY CREATE, {self.data.close[0]}')
                self.order = self.buy()
        else:
            if self.ma[0] < self.ma[-1] and self.ma[-1] > self.ma[-2]:
                self.log(f'SELL CREATE, {self.data.close[0]}')
                self.order = self.sell()

if __name__ == '__main__':
    cerebro = bt.Cerebro()

    # 添加数据源
    symbol = "510050"
    adjust = "hfq"
    start_date = datetime.datetime(2023, 5, 1)
    end_date = datetime.datetime(2024, 6, 28)
    etf = etf_daily(symbol,  adjust)
    data = etf.get_data()
    data = MyStockData(dataname=data, fromdate=start_date, todate=end_date) 
    cerebro.adddata(data)

    # Add strategy
    cerebro.addstrategy(MovingAverageTurnStrategy)

    # Set initial cash
    cerebro.broker.setcash(100000.0)

    # Run the strategy
    cerebro.run()

    # Plot the result
    cerebro.plot()
    
    
    